# -*- coding: UTF-8 -*-
'''
Created on 2020年8月25日

@author: Administrator
'''
import tensorflow as tf
import falcon
from wsgiref import simple_server

import yaml
import numpy as np
import matplotlib.pyplot as plt

import IPython.display as ipd

from tensorflow_tts.inference import AutoConfig
from tensorflow_tts.inference import TFAutoModel
from tensorflow_tts.inference import AutoProcessor
import re
import io
from scipy.io import wavfile
import os
os.environ["CUDA_VISIBLE_DEVICES"]="1" #字段GPU

tacotron2_config = AutoConfig.from_pretrained('../examples/tacotron2/conf/tacotron2.baker.v1.yaml')
tacotron2 = TFAutoModel.from_pretrained(
    config=tacotron2_config,
    pretrained_path="tacotron2-100k.h5",
    training=False, 
    name="tacotron2"
)

fastspeech2_config = AutoConfig.from_pretrained('../examples/fastspeech2/conf/fastspeech2.baker.v2.yaml')
fastspeech2 = TFAutoModel.from_pretrained(
    config=fastspeech2_config,
    pretrained_path="fastspeech2-200k.h5",
    name="fastspeech2"
)

mb_melgan_config = AutoConfig.from_pretrained('../examples/multiband_melgan/conf/multiband_melgan.baker.v1.yaml')
mb_melgan = TFAutoModel.from_pretrained(
    config=mb_melgan_config,
    pretrained_path="mb.melgan-920k.h5",
    name="mb_melgan"
)

processor = AutoProcessor.from_pretrained(pretrained_path="baker_mapper.json")

def do_synthesis(input_text, text2mel_model, vocoder_model, text2mel_name, vocoder_name):
  input_ids = processor.text_to_sequence(input_text, inference=True)

  # text2mel part
  if text2mel_name == "TACOTRON":
    _, mel_outputs, stop_token_prediction, alignment_history = text2mel_model.inference(
        tf.expand_dims(tf.convert_to_tensor(input_ids, dtype=tf.int32), 0),
        tf.convert_to_tensor([len(input_ids)], tf.int32),
        tf.convert_to_tensor([0], dtype=tf.int32)
    )
  elif text2mel_name == "FASTSPEECH2":
    mel_before, mel_outputs, duration_outputs, _, _ = text2mel_model.inference(
        tf.expand_dims(tf.convert_to_tensor(input_ids, dtype=tf.int32), 0),
        speaker_ids=tf.convert_to_tensor([0], dtype=tf.int32),
        speed_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
        f0_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
        energy_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
    )
  elif text2mel_name == "MB_MELGAN":
    mel_before, mel_outputs, duration_outputs, _, _ = text2mel_model.inference(
        tf.expand_dims(tf.convert_to_tensor(input_ids, dtype=tf.int32), 0),
        speaker_ids=tf.convert_to_tensor([0], dtype=tf.int32),
        speed_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
        f0_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
        energy_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
    )
  else:
    raise ValueError("Only TACOTRON, FASTSPEECH2 are supported on text2mel_name")

  # vocoder part
  if vocoder_name == "MB-MELGAN":
    # tacotron-2 generate noise in the end symtematic, let remove it :v.
    if text2mel_name == "TACOTRON":
      remove_end = 1024
    else:
      remove_end = 1
    audio = vocoder_model.inference(mel_outputs)[0, :-remove_end, 0]
  else:
    raise ValueError("Only MB_MELGAN are supported on vocoder_name")

  if text2mel_name == "TACOTRON":
    return mel_outputs.numpy(), alignment_history.numpy(), audio.numpy()
  else:
    return mel_outputs.numpy(), audio.numpy()

# setup window for tacotron2 if you want to try
tacotron2.setup_window(win_front=5, win_back=5)

html_body = '''<html><title>tensorflowTTS Demo</title><meta charset='utf-8'>
<style>
body {padding: 16px; font-family: sans-serif; font-size: 14px; color: #444}
input {font-size: 14px; padding: 8px 12px; outline: none; border: 1px solid #ddd}
input:focus {box-shadow: 0 1px 2px rgba(0,0,0,.15)}
p {padding: 12px}
button {background: #28d; padding: 9px 14px; margin-left: 8px; border: none; outline: none;
        color: #fff; font-size: 14px; border-radius: 4px; cursor: pointer;}
button:hover {box-shadow: 0 1px 2px rgba(0,0,0,.15); opacity: 0.9;}
button:active {background: #29f;}
button[disabled] {opacity: 0.4; cursor: default}
</style>
<body>
<form>
  <input id="text" type="text" size="40" placeholder="请输入文字">
  <button id="button" name="synthesize">合成</button>
  <div>模型:
        <label><input type="radio" name="model" checked value="TACOTRON">TACOTRON</label>
        <label><input type="radio" name="model" value="FASTSPEECH2">FASTSPEECH2</label>
    </div>
</form>
<p id="message"></p>
<audio id="audio" controls autoplay hidden></audio>
<script>
function q(selector) {return document.querySelector(selector)}
q('#text').focus()
q('#button').addEventListener('click', function(e) {
  text = q('#text').value.trim()
  model = q("input[name='model']:checked").value.trim()
  if (!model) {
      model='TACOTRON'
  }
  if (text) {
    q('#message').textContent = '合成中...'
    q('#button').disabled = true
    q('#audio').hidden = true
    synthesize(text,model)
  }
  e.preventDefault()
  return false
})
function synthesize(text,model) {
  fetch('/synthesize?text=' + encodeURIComponent(text)+'&model='+model, {cache: 'no-cache'})
    .then(function(res) {
      if (!res.ok) throw Error(res.statusText)
      return res.blob()
    }).then(function(blob) {
      q('#message').textContent = ''
      q('#button').disabled = false
      q('#audio').src = URL.createObjectURL(blob)
      q('#audio').hidden = false
    }).catch(function(err) {
      q('#message').textContent = '出错: ' + err.message
      q('#button').disabled = false
    })
}
</script></body></html>
'''
def find_chinese(str_data):
    '''去除非中文'''
    pattern = re.compile(r'[^\u4e00-\u9fa5]')
    chinese = re.sub(pattern, '#3', str_data)
    return chinese

def p(input,model):
    input_text=find_chinese(input)
    if model=='TACOTRON':
        mels, alignment_history, audios = do_synthesis(input_text, tacotron2, mb_melgan, "TACOTRON", "MB-MELGAN")
    else:
        mels, audios = do_synthesis(input_text, fastspeech2, mb_melgan, "FASTSPEECH2", "MB-MELGAN")    
        
    out = io.BytesIO()
    audios *= 32767 / max(0.01, np.max(np.abs(audios)))
    wavfile.write(out, 24000, audios.astype(np.int16))
    return out.getvalue()

    
class Res:
    def on_get(self,req,res):
        res.body = html_body
        res.content_type = "text/html"

class Syn:
    def on_get(self,req,res):
        if not req.params.get('text'):
            raise falcon.HTTPBadRequest()
        text=req.params.get('text')
        model=req.params.get('model')
        res.data = p(text,model)
        res.content_type = "audio/wav"		

api = falcon.API()
api.add_route("/",Res())
api.add_route("/synthesize",Syn())
print("host:{},port:{}".format('localhost',int(9000)))
simple_server.make_server('localhost',int(9000),api).serve_forever()






